package com.ideaaedi.component.dump;

import lombok.extern.slf4j.Slf4j;
import net.bytebuddy.agent.ByteBuddyAgent;
import org.apache.commons.lang3.StringUtils;

import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.Instrumentation;
import java.lang.instrument.UnmodifiableClassException;
import java.security.ProtectionDomain;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;

/**
 * 使用ClassFileTransformer dump当前JVM中的class
 * <br><br/>
 * 注：与{@link NonExitDumpClassExecutor}一样，使用此工具类dump户，并不会退出当前进程。
 * <br><br/>
 * 特别说明：因为Instrumentation触发的是一个ClassFileTransformer链，可能这条ClassFileTransformer链上有很多ClassFileTransformer，
 *          当{@link Instrumentation#addTransformer(ClassFileTransformer)}时，只是把你的这个ClassFileTransformer添加至这条链的最后，
 *          所以当执行到你的ClassFileTransformer实例时，{@link ClassFileTransformer#transform}入参传过来的字节码，可能已经被前面的
 *          ClassFileTransformer做了手脚。
 * <br/>
 * <br/>
 * 所以，使用ClassFileTransformer dump的class并不一定是JVM中正运行着的class，如果你想要获取JVM上正运行着的class的话，
 * 你可以考虑使用{@link NonExitDumpClassExecutor}或{@link ExitDumpClassExecutor}，当然，它们也有他们的不足之处。
 * <br/>
 * <br/>
 * 最后，虽然ClassFileTransformer有这个机制问题，但是绝不部分情况下，前面的ClassFileTransformer都不会篡改class字节码。
 *    因此，一定程度上，我们可以放心使用ExitDumpClassExecutor。
 *
 * @author JustryDeng
 * @since 2021/9/27 22:01:28
 */
@Slf4j
public class NonExitClassFileTransformerExecutor implements ClassFileTransformer {
    
    /** 要dump的class的全类名前缀匹配集合 */
    private final Set<String> INCLUDE_PREFIX_SET = Collections.synchronizedSet(new HashSet<>());
    
    /** 明确排除dump的class的全类名前缀匹配集合 */
    private final Set<String> EXCLUDE_PREFIX_SET = Collections.synchronizedSet(new HashSet<>());
    
    /**
     * 是否以全类名作为{@link NonExitClassFileTransformerExecutor#dumpClassMap}的key
     * <br><br>
     * 假设现有类com.example.HelloWorld，那么：
     * <ul>
     *     <li>为true时：key形如com.example.HelloWorld</li>
     *     <li>为false时：key形如com/example/HelloWorld</li>
     * </ul>
     */
    public final boolean CLASS_LONG_NAME_AS_KEY;
    
    /**
     * dump的class
     * <pre>
     *     k - 类名, 格式详见{@link NonExitClassFileTransformerExecutor#CLASS_LONG_NAME_AS_KEY}
     *     v -
     *        k-加载当前class的类加载器信息，
     *        v-class字节码
     * </pre>
     */
    private final Map<String, Map<String, byte[]>> dumpClassMap = new ConcurrentHashMap<>();
    
    public static NonExitClassFileTransformerExecutor create(String includePrefixes) {
        return new NonExitClassFileTransformerExecutor(includePrefixes, null, ExitDumpClassExecutor.DEFAULT_KEEP_LONG_NAME);
    }
    
    public static NonExitClassFileTransformerExecutor create(String includePrefixes, String excludePrefixes) {
        return new NonExitClassFileTransformerExecutor(includePrefixes, excludePrefixes, ExitDumpClassExecutor.DEFAULT_KEEP_LONG_NAME);
    }

    
    public static NonExitClassFileTransformerExecutor create(String includePrefixes, boolean keepLongName) {
        return new NonExitClassFileTransformerExecutor(includePrefixes, null, keepLongName);
    }
    
    public static NonExitClassFileTransformerExecutor create(String includePrefixes, String excludePrefixes, boolean keepLongName) {
        return new NonExitClassFileTransformerExecutor(includePrefixes, excludePrefixes, keepLongName);
    }
    
    private NonExitClassFileTransformerExecutor(String includePrefixes, String excludePrefixes, boolean keepLongName) {
        // INCLUDE_PREFIX_SET
        if (StringUtils.isBlank(includePrefixes)) {
            throw new IllegalArgumentException("includePrefixes cannot be blank.");
        }
        this.INCLUDE_PREFIX_SET.addAll(Arrays.stream(includePrefixes.split(","))
                .map(String::trim)
                .filter(StringUtils::isNotBlank)
                .collect(Collectors.toSet()));
        
        // EXCLUDE_PREFIX_SET
        if (StringUtils.isNotBlank(excludePrefixes)) {
            this.EXCLUDE_PREFIX_SET.addAll(Arrays.stream(excludePrefixes.split(","))
                    .map(String::trim)
                    .filter(StringUtils::isNotBlank)
                    .collect(Collectors.toSet()));
        }
        // KEEP_LONG_NAME
        this.CLASS_LONG_NAME_AS_KEY = keepLongName;
        log.info("INCLUDE_PREFIX_SET -> {}, EXCLUDE_PREFIX_SET -> {}, KEEP_LONG_NAME -> {}", INCLUDE_PREFIX_SET, EXCLUDE_PREFIX_SET, CLASS_LONG_NAME_AS_KEY);
    }
    
    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, 
                            ProtectionDomain protectionDomain, byte[] classfileBuffer) {
        // className形如 com/ideaaedi/component/dump/NonExitClassFileTransformerExecutor$$Lambda$274
        String classLongName = className.replace("/", ".");
        // 再判断一下是否应该dump, 因为有一个ClassFileTransformer链，当执行到当前的ClassFileTransformer时，你并不能确定是否前面有什么骚操作产生了新的Class
        if (!ifDump(classLongName)) {
            return classfileBuffer;
        }
        String key;
        if (CLASS_LONG_NAME_AS_KEY) {
            key = classLongName;
        } else {
            key = className;
        }
        String classLoaderInfo = loader.getClass().getTypeName() + "@" + loader.hashCode();
        // 已处理了的不再重复处理
        if (dumpClassMap.containsKey(key) && dumpClassMap.get(key).containsKey(classLoaderInfo)) {
            return classfileBuffer;
        }
        if (dumpClassMap.containsKey(key)) {
            dumpClassMap.get(key).put(classLoaderInfo, classfileBuffer);
        } else {
            Map<String, byte[]> innerMap = new ConcurrentHashMap<>(4);
            innerMap.put(classLoaderInfo, classfileBuffer);
            dumpClassMap.put(key, innerMap);
        }
        return classfileBuffer;
    }
    
    /**
     * dump 指定JVM中的class
     *
     * @return key - 类名,   value - class字节码
     */
    public Map<String, Map<String, byte[]>> exec() {
        dumpClassMap.clear();
        Instrumentation instrumentation = ByteBuddyAgent.install();
        instrumentation.addTransformer(this, true);
        try {
            //noinspection rawtypes
            Class[] allLoadedClasses = instrumentation.getAllLoadedClasses();
            //noinspection rawtypes
            Class[] needDumpClasses = Arrays.stream(allLoadedClasses).filter(klass -> {
                String classLoneName = klass.getTypeName();
                return ifDump(classLoneName);
            }).toArray(Class[]::new);
            if (needDumpClasses.length > 0) {
                // 再次进行transform， 使触发回调进而调用ClassFileTransformer链
                instrumentation.retransformClasses(needDumpClasses);
            }
        } catch (UnmodifiableClassException e) {
            log.error("ClassFileTransformerExecutor dump class exception", e);
        } finally {
            instrumentation.removeTransformer(this);
        }
        return dumpClassMap;
    }
    
    /**
     * 是否应该dump
     *
     * @param classLoneName
     *            全类名，即{@link Class#getTypeName()}
     * @return  是否应该dump
     */
    private boolean ifDump(String classLoneName) {
        for (String excludePrefix : EXCLUDE_PREFIX_SET) {
            if (classLoneName.startsWith(excludePrefix)) {
                return false;
            }
        }
        for (String includePrefix : INCLUDE_PREFIX_SET) {
            if (classLoneName.startsWith(includePrefix)) {
                return true;
            }
        }
        return false;
    }
}
